import numpy as np
from numba import njit
from numba_stats import norm
"""
This module defines a collection of drift functions for the simulators in the package.
"""


def constant(t=np.arange(0, 20, 0.1)):
    """constant drift function

    Arguments
    ---------
        t (_type_, optional): _description_. Defaults to np.arange(0, 20, 0.1).

    Returns
    -------
        np.array: Array of drift values, same length as t

    """
    return np.zeros(t.shape[0])


def gamma_drift(t=np.arange(0, 20, 0.1), shape=2, scale=0.01, c=1.5):
    """Drift function that follows a scaled gamma distribution

    Arguments
    ---------
        t: np.ndarray <default=np.arange(0, 20, 0.1)>
            Timepoints at which to evaluate the drift.
            Usually np.arange() of some sort.
        shape: float <default=2>
            Shape parameter of the gamma distribution
        scale: float <default=0.01>
            Scale parameter of the gamma distribution
        c: float <default=1.5>
            Scalar parameter that scales the peak of
            the gamma distribution.
            (Note this function follows a gamma distribution
            but does not integrate to 1)

    Return
    ------
        np.ndarray
            The gamma drift evaluated at the supplied timepoints t.

    """

    num_ = np.power(t, shape - 1) * np.exp(np.divide(-t, scale))
    div_ = (
        np.power(shape - 1, shape - 1) * np.power(scale, shape - 1) *
        np.exp(-(shape - 1))
    )
    return c * np.divide(num_, div_)


def ds_support_analytic(t=np.arange(0, 10, 0.001), init_p=0, fix_point=1, slope=2):
    """Solution to differential equation of the form:
       x' = slope*(fix_point - x),
       with initial condition init_p.
       The solution takes the form:
       (init_p - fix_point) * exp(-slope * t) + fix_point

    Arguments
    ---------
        t: np.ndarray <default=np.arange(0, 20, 0.1)>
            Timepoints at which to evaluate the drift. Usually np.arange() of some sort.
        init_p: float <default=0>
            Initial condition of dynamical system
        fix_point: float <default=1>
            Fixed point of dynamical system
        slope: float <default=0.01>
            Coefficient in exponent of the solution.
    Return
    ------
    np.ndarray
         The gamma drift evaluated at the supplied timepoints t.

    """

    return (init_p - fix_point) * np.exp(-(slope * t)) + fix_point


def ds_conflict_drift(
    t=np.arange(0, 10, 0.001),
    init_p_t=0,
    init_p_d=0,
    slope_t=1,
    slope_d=1,
    fixed_p_t=1,
    coherence_t=1.5,
    coherence_d=1.5
):
    """This drift is inspired by a conflict task which
        involves a target and a distractor stimuli both presented
        simultaneously.
        Two drift timecourses are linearly combined weighted
        by the coherence in the respective target and distractor stimuli.
        Each timecourse follows a dynamical system as described
        in the ds_support_analytic() function.

    Arguments
    ---------
        t: np.ndarray <default=np.arange(0, 20, 0.1)>
            Timepoints at which to evaluate the drift.
            Usually np.arange() of some sort.
        init_p_t: float <default=0>
            Initial condition of target drift timecourse
        init_p_d: float <default=0>
            Initial condition of distractor drift timecourse
        slope_t: float <default=1>
            Slope parameter for target drift timecourse
        slope_d: float <default=1>
            Slope parameter for distractor drift timecourse
        fixed_p_t: float <default=1>
            Fixed point for target drift timecourse
        coherence_t: float <default=1.0>
            Coefficient for the target drift timecourse
        coherence_d: float <default=-1.0>
            Coefficient for the distractor drift timecourse
    Return
    ------
    np.ndarray
        The full drift timecourse evaluated at the supplied timepoints t.
    """

    w_t = ds_support_analytic(
        t=t, init_p=init_p_t, fix_point=fixed_p_t, slope=slope_t
    )

    w_d = ds_support_analytic(t=t, init_p=init_p_d, fix_point=0, slope=slope_d)

    v_t = (w_t * coherence_t) + (w_d * coherence_d)

    return v_t  # , w_t, w_d


@njit(fastmath=True)
def SSP_drift(t: float, p: float, sda: float, rd: float) -> np.array:
    """Drift function of Spotlight Shrinkage model for Diffusion model in conflict tasks by Evans et al. (2020)

    Args:
        t (float): time of sequential sampling in seconds
        p (float): point of drift rate
        sda (float): inital deviation of drift rate
        rd (float): reduction rate of inital deviation

    Returns:
        np.array: np.array include scalar drift rate
    """
    # calculate current sd of spotlight
    sd_t = sda - (rd * t)
    sd_t = max(sd_t, 0.001)

    x = np.array([0.5], dtype=np.float32)
    # find area of spotlight over target and flanker
    # NOTE numba_stats 1.7.0: norm.cdf only support x is np.array, mu and sigma is float
    a_target = norm.cdf(x, 0.0, sd_t) - norm.cdf(-x, 0.0, sd_t)
    a_flanker = 1 - a_target

    # current drift rate
    drift = p * (a_target - a_flanker)
    # drift = 2 * p * (a_target - 0.5)

    return drift


@njit(fastmath=True)
def DMC_drift(t: float, vc: float, peak: float, shape: float, tau: float) -> float:
    """Drift function for Diffusion model in conflict tasks that follows a gamma function by Evans et al. (2020)

    Arguments
    ---------
        t: float
            Timepoints at which to evaluate the drift.
        vc: float
            The drift for control process. 
        shape: float 
            Shape parameter of the gamma distribution
        peak: float
            Scalar parameter that scales the peak of
            the gamma distribution.
            (Note this function follows a gamma distribution
            but does not integrate to 1)
        tau: float
            tau is the characteristic time.

    Return
    ------
        float
            The gamma drift evaluated at the supplied timepoints t without congruency condition.
    """

    t = max(t, 0.01)
    term1 = peak * np.exp(-t / tau)
    term2 = np.power((t * np.e) / ((shape - 1) * tau), (shape - 1))
    term3 = ((shape - 1) / t) - (1 / tau)
    va = term1 * term2 * term3

    drift = va + vc
    return drift

@njit
def rDMC_drift(t: float, k: float, A_0: float, d_a: float, d_c: float) -> float:
    """Lee, P., & Sewell, D. (2024). A revised diffusion model for conflict tasks. Psychonomic Bulletin & Review, 31(1), 1–31. https://doi.org/10.3758/s13423-023-02288-0

    $$
    \\begin{aligned}
    &w_a(t)=A_0e^{-kt}\\
    &w_c(t)=1-w_a(t)\\
    &v(t)=w_a(t)d_a+w_c(t)d_c\\
    \\end{aligned}
    $$

    Note: In incongruent condition, d_a to -d_a

    Return
    ------
        float
            The gamma drift evaluated at the supplied timepoints t without congruency condition.
    """

    # automatic channel activation weight
    w_a = A_0 * np.exp(-k * t)
    # control channel activation weight
    # w_c = 1 - w_a

    # drift = w_a * d_a + w_c * d_c
    drift = w_a * (d_a - d_c) + d_c

    return drift

# @njit(fastmath=True)
# def DSTP_drift(
#     X_ss: float,
#     mu: float,
#     # vta: float,
#     # vfl: float,
#     vss: float,
#     vp2: float,
#     ass: float,
#     # zss: float,
#     congruency=1,
#     delta_t=0.001,
#     sqrt_st=np.sqrt(0.001)
# ):
#     """Drift function for dual-stage two-phase modell in conflict tasks that follows a gamma function by Evans et al. (2020). The code is transformed from Luo, J., Yang, M., & Wang, L. (2022). Learned irrelevant stimulus-response associations and proportion congruency effect: A diffusion model account. Journal of Experimental Psychology: Learning, Memory, and Cognition. https://doi.org/10.1037/xlm0001158

#     Arguments
#     ---------
#         t: foat
#             Represents the timepoints at which to evaluate the drift.
#         vta: float
#             Drift for target in phase 1.
#         vfl: float
#             Drift for flanker in phase 1.
#         vss: float
#             Drift for stimulus selection.
#         ass: float
#             Boundary for stimulus selection.
#         zss: float <default=0.5>
#             Start point for stimulus selection.
#             Indicates the presence or absence of selection bias.
#             Defaults to 0.5, which means no selection bias.
#         congruency: int
#             Congruency condition with flanker and target.
#             1 means congruent, -1 means incongruent.
#         delta_t: float
#             Time step for diffusion process.
#         sqrt_st: float
#             Scale factor for diffusion process.

#     Return
#     ------
#         np.ndarray
#             The full drift timecourse evaluated at the supplied timepoints t.
#     """

#     # initiate the X of start point for stimulus selection
#     # X_ss = zss * ass

#     # Drift rate in the first phase
#     # mu = vta + congruency * vfl

#     # Stimulus selection
#     # X_ss += vss * delta_t + np.random.normal() * sqrt_st
#     X_ss += vss
#     # Drift rate in the second phase
#     if X_ss >= ass:
#         drift = vp2
#     elif X_ss <= 0 and congruency == 0:
#         drift = vp2
#     # elif X_ss <= 0 and congruency == -1:
#     #     drift = -(vp2)
#     return drift

# @njit(fastmath=True)
# def DSTP_drift(
#     t: float,
#     vta: float,
#     vfl: float,
#     vss: float,
#     vp2: float,
#     ass: float,
#     zss: float,
#     congruency=1,
#     delta_t=0.001,
#     sqrt_st=np.sqrt(0.001)
# ):
#     """Drift function for dual-stage two-phase modell in conflict tasks that follows a gamma function by Evans et al. (2020). The code is transformed from Luo, J., Yang, M., & Wang, L. (2022). Learned irrelevant stimulus-response associations and proportion congruency effect: A diffusion model account. Journal of Experimental Psychology: Learning, Memory, and Cognition. https://doi.org/10.1037/xlm0001158

#     Arguments
#     ---------
#         t: np.ndarray <default=np.arange(0, 20, 0.001)>
#             Drift timecourse.
#             Represents the timepoints at which to evaluate the drift.
#         vta: float
#             Drift for target in phase 1.
#         vfl: float
#             Drift for flanker in phase 1.
#         vss: float
#             Drift for stimulus selection.
#         ass: float
#             Boundary for stimulus selection.
#         zss: float <default=0.5>
#             Start point for stimulus selection.
#             Indicates the presence or absence of selection bias.
#             Defaults to 0.5, which means no selection bias.
#         congruency: int
#             Congruency condition with flanker and target.
#             1 means congruent, -1 means incongruent.
#         delta_t: float
#             Time step for diffusion process.
#         sqrt_st: float
#             Scale factor for diffusion process.

#     Return
#     ------
#         np.ndarray
#             The full drift timecourse evaluated at the supplied timepoints t.
#     """

#     # initiate the X of start point for stimulus selection
#     X_ss = zss * ass

#     # Drift rate in the first phase
#     congruency = np.min([congruency, 0])
#     mu = vta + congruency * vfl

#     t = t.shape[0]
#     mu_view = np.full(t, mu)
#     for step in range(t):
#         # Stimulus selection
#         X_ss += vss * delta_t + np.random.normal() * sqrt_st
#         # Drift rate in the second phase
#         if X_ss >= ass:
#             mu_view[step:] = vp2
#             break
#         elif X_ss <= 0 and congruency == 0:
#             mu_view[step:] = vp2
#             break
#         elif X_ss <= 0 and congruency == -1:
#             mu_view[step:] = -(vp2)
#             break
#     return mu_view
